#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
将预测CSV文件中的UTC时间转换为北京时间（UTC+8）
使用方法: python convert_utc_to_beijing.py <input_csv_file>
"""

import pandas as pd
import argparse
from datetime import timedelta
import os

def convert_utc_to_beijing(input_file):
    """
    将CSV文件中的UTC时间转换为北京时间
    
    Args:
        input_file (str): 输入CSV文件路径
        
    Returns:
        str: 输出文件路径
    """
    # 读取CSV文件
    df = pd.read_csv(input_file)
    
    # 确保timestamp列存在
    if 'timestamp' not in df.columns:
        raise ValueError("CSV文件中必须包含'timestamp'列")
    
    # 将timestamp列转换为datetime对象（带UTC时区）
    df['timestamp'] = pd.to_datetime(df['timestamp'], utc=True)
    
    # 直接转换为北京时间时区 (UTC+8)
    df['timestamp'] = df['timestamp'].dt.tz_convert('Asia/Shanghai')
    
    # 生成输出文件名
    base_name = os.path.splitext(input_file)[0]
    output_file = f"{base_name}_beijing.csv"
    
    # 保存为新的CSV文件
    df.to_csv(output_file, index=False)
    
    print(f"转换完成！输出文件: {output_file}")
    print(f"时间范围: {df['timestamp'].min()} 到 {df['timestamp'].max()}")
    
    return output_file

def main():
    parser = argparse.ArgumentParser(description='将UTC时间转换为北京时间')
    parser.add_argument('input_file', help='输入CSV文件路径')
    
    args = parser.parse_args()
    
    try:
        convert_utc_to_beijing(args.input_file)
    except Exception as e:
        print(f"错误: {e}")
        print("请确保文件存在且包含'timestamp'列")

if __name__ == "__main__":
    main()
